import torch
from torchvision import models
import logging
from logging.handlers import RotatingFileHandler
from flask import Flask, request, jsonify
from utils.prepressing import preprocess_image
from utils.postpressing import postprocess_output
from utils.files_handling import save_result_to_file
import datetime
import os

# 创建日志记录器
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# 创建文件处理器，设置日志文件的相关参数，如最大文件大小、备份文件数量等
file_handler = RotatingFileHandler('logs/app.log', maxBytes=1024 * 1024 * 10, backupCount=5)
# 创建日志格式器
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler.setFormatter(formatter)

# 将文件处理器添加到日志记录器
logger.addHandler(file_handler)
# 定义加载模型的函数
def load_model():
    model = models.resnet18(pretrained=False)  # 这里以ResNet18为例，替换为你实际的模型结构定义
    #model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
    model.eval()
    return model

# 加载模型，假设模型文件在'model/your_trained_model.pth'
model = load_model()#'model/your_trained_model.pth')

app = Flask(__name__)

@app.route('/predict', methods=['POST'])
def predict():
    data = request.get_json()  # 获取POST请求中的JSON数据
    logger.info(f"Received request data: {data}")
    if data is None:
        logger.error("Invalid input format, expecting JSON data")
        return jsonify({"error": "Invalid input format, expecting JSON data"}), 400
    # 后续在这里进行输入前处理、模型推断等操作，暂先返回一个占位消息
    input_data = preprocess_image(data['image'])  # 假设请求的JSON数据中有'image'字段存放图像数据
    with torch.no_grad():  # 如果是PyTorch模型，关闭梯度计算以节省内存并加速推断（对于TensorFlow通常不需要这步）
        output = model(torch.from_numpy(input_data).unsqueeze(0).float())  # 增加批次维度并转换数据类型，按实际模型输入要求调整
    result = postprocess_output(output)
    # 定义保存文件的路径，这里假设保存到项目根目录下的'results'文件夹中的以当前时间命名的文本文件（示例，可按实际改）
    file_name = datetime.datetime.now().strftime('%Y%m%d%H%M%S') + '.txt'
    file_path = os.path.join('results', file_name)
    save_result_to_file(result, file_path)
    # 后续进行输出后处理等操作
    return jsonify({"result": result})

from flask_profiler import Profiler

app = Flask(__name__)
app.config["DEBUG"] = True

# 配置Profiler，指定要监控的函数等（这里简单配置全部路由，按实际需求调整）
app.config["flask_profiler"] = {
    "enabled": True,
    "storage": {
        "engine": "sqlite"  # 可以选择存储监控数据的后端，如sqlite、mysql等
    },
    "basicAuth":{
        "enabled": True,
        "username": "admin",
        "password": "admin"
    },
    "ignore": [
        "^/static/.*"
    ]
}
Profiler(app)

# 下面是前面定义的路由等代码，此处省略重复内容

if __name__ == '__main__':
    app.run(debug=True)